#  Copyright (c) 2025 Hovi, 回血池. All rights reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
import time
import os


def create_data_directories():
    """创建数据存储目录"""
    directories = ['data/stocks', 'data/indices']
    for directory in directories:
        if not os.path.exists(directory):
            os.makedirs(directory)
            print(f"创建目录: {directory}")


def download_stock_data():
    """下载股票数据"""
    try:
        # 获取所有A股列表
        stock_list = ak.stock_zh_a_spot()
        total_stocks = len(stock_list)

        for idx, row in stock_list.iterrows():
            try:
                raw_code = row['代码']  # 可能形如 'sh600000' 或 'sz000001'
                name = row['名称']

                # 提取数字部分作为股票代码
                code = ''.join(filter(str.isdigit, raw_code))

                # 验证是否为有效的6位股票代码
                if not (code.isdigit() and len(code) == 6):
                    print(f"跳过无效股票代码: {raw_code}")
                    continue

                filename = f"data/stocks/{code}.csv"

                # 如果文件已存在则跳过
                if os.path.exists(filename):
                    print(f"已存在，跳过: {code}_{name}")
                    continue

                # 获取该股票的历史数据
                df = ak.stock_zh_a_hist(
                    symbol=code,
                    period="daily",
                    start_date="20190101",
                    end_date="20240108",
                    adjust="qfq"
                )

                # 保存数据
                df.to_csv(filename, index=False, encoding='utf-8-sig')
                print(f"已下载 [{idx + 1}/{total_stocks}]: {code}_{name}")

                time.sleep(1)

            except Exception as e:
                print(f"下载失败 {raw_code}_{name}: {e}")
                continue

    except Exception as e:
        print(f"获取股票列表失败: {e}")


def download_index_data():
    """下载主要指数数据"""
    indices = {
        'sh000001': '上证指数',  # 上交所
        'sz399001': '深证成指',  # 深交所
        'bj899050': '北证50',  # 北交所
        'sz399006': '创业板指',  # 深交所
        'sh000688': '科创50',  # 上交所
        'sz399550': '深主板50',  # 深交所
        'sh000300': '沪深300',  # 上交所
        'sh000016': '上证50',  # 上交所
        'sz399050': '深证50',  # 深交所
        'sz399101': '中小100',  # 深交所
        'csi000050': '中证A50',  # 中证指数
        'csi000903': '中证A100',  # 中证指数
        'csi000905': '中证500',  # 中证指数
        'csi000200': '中证200',  # 中证指数
        'csi000906': '中证800',  # 中证指数
        'csi000852': '中证1000',  # 中证指数
        'csi932000': '中证2000',  # 中证指数
        'csi000985': '中证全指',  # 中证指数
        'sh000010': '上证180',  # 上交所
        'sh000009': '上证380',  # 上交所
        'sh000903': '上证100',  # 上交所
        'sh000150': '上证150'  # 上交所
    }

    for code, name in indices.items():
        try:
            filename = f"data/indices/{code}.csv"

            if os.path.exists(filename):
                print(f"已存在，跳过: {code}_{name}")
                continue

            # 获取指数数据
            df = ak.stock_zh_index_daily(symbol=code)


            # stock_zh_index_daily_em_df = ak.stock_zh_index_daily_em(symbol=code)
            # 保存数据
            df.to_csv(filename, index=False, encoding='utf-8-sig')
            print(f"已下载: {code}_{name}")

            time.sleep(1)

        except Exception as e:
            print(f"下载失败 {code}_{name}: {e}")
            continue


def main():
    print("开始下载A股数据...")
    create_data_directories()

    print("\n下载股票数据...")
    download_stock_data()

    print("\n下载指数数据...")
    download_index_data()

    print("\n数据下载完成！")


if __name__ == "__main__":
    main()
